﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Bioinformatics.Types
{
    public class Particle
    {
        public static ParticleComparer Comparer = new ParticleComparer();
        private static Random _random = new Random();
        public static float[] BranchLengthGlobalBest { get; set; }

        public Phylogeny Phylogeny { get; set; }
        public float[] BranchLengths { get; private set; }
        public float[] BranchDeltas { get; set; }
        public float[] BranchLengthsBest { get; set; }
        public float[] BranchLengthsNeighbourhoodBest { get; set; }
        public int[] Neighbourhood { get; private set; }
        public float Likelihood { get; private set; }
        public float LikelihoodBest { get; private set; }
        

        public Particle(Phylogeny phylogeny, float branchLengthMax, float branchLengthMin, float branchDeltaMax, float branchDeltaMin, int swarmSize, int neighbourhoodSize)
        {
            Phylogeny = phylogeny;

            var range = branchLengthMax - branchLengthMin;

            BranchLengths = new float[Phylogeny.PhylogenyNodes.Length];
            for (var i = 0; i < BranchLengths.Length; i++)
                BranchLengths[i] = branchLengthMin + range*(float) _random.NextDouble();
            BranchLengthsBest = (float[]) BranchLengths.Clone();
            LikelihoodBest = float.MinValue;
               

            range = branchDeltaMax - branchDeltaMin;
            BranchDeltas = new float[BranchLengths.Length];
            for (var i = 0; i < BranchDeltas.Length; i++)
                BranchDeltas[i] = branchDeltaMin + range * (float)_random.NextDouble();

            Neighbourhood = new int[neighbourhoodSize];
            for (var i = 0; i < neighbourhoodSize; i++)
                Neighbourhood[i] = _random.Next(swarmSize);
        }

        public void CalculateLikelihood()
        {
            //variable to hold total Probability
            Likelihood = 0.0f;

            for (var site = 0; site < Phylogeny.Sequences.Count; site++)
                foreach (var character in Phylogeny.StationaryProbabilities.Keys)
                    Likelihood +=
                        likelihoodRecursion(0, character, site)*Phylogeny.StationaryProbabilities[character];
            
            //update best solution seen so far
            if (Likelihood > LikelihoodBest)
            {
                BranchLengthsBest = (float[]) BranchLengths.Clone();
                LikelihoodBest = Likelihood;
            }
        }

        private float likelihoodRecursion(int currentIdx, char parentState, int sequenceStateIdx)
        {
            //check to see if the current node is a tip
            if(Phylogeny.PhylogenyNodes[currentIdx].Value != string.Empty)
                return Phylogeny.Sequences[Phylogeny.PhylogenyNodes[currentIdx].Value][sequenceStateIdx] == parentState ? 1.0f : 0.0f;

            //variables to hold left and right subtree likelihoods
            var leftSubtreeLikelihood = 0.0f;
            var rightSubtreeLikelihood = 0.0f;

            //summate over all possible states
            foreach (var state in Phylogeny.StationaryProbabilities.Keys)
            {
                //calculate probability of branch configuration multiplied by subtree Probability for both subtrees
                leftSubtreeLikelihood +=
                    Probability(parentState, state, BranchLengths[currentIdx]) *
                    likelihoodRecursion(Phylogeny.PhylogenyNodes[currentIdx].LeftChild, state, sequenceStateIdx);
                rightSubtreeLikelihood +=
                    Probability(parentState, state, BranchLengths[currentIdx]) *
                    likelihoodRecursion(Phylogeny.PhylogenyNodes[currentIdx].RightChild, state, sequenceStateIdx);
            }

            //return the product of the two
            return leftSubtreeLikelihood * rightSubtreeLikelihood;
        }

        private static float Probability(char stateI, char stateJ, float length)
        {
            var exp = (float)Math.Exp(-4.0f * 0.1f * length);
            if (stateI == stateJ)
                return 0.25f + 0.75f * exp;
            return 0.25f - 0.25f * exp;
        }

        public void UpdateBranchLengths()
        {
            for (var i = 0; i < BranchLengths.Length; i++)
                BranchLengths[i] += BranchDeltas[i];
        }

        public void AdjustVelocities()
        {
            for (var b = 0; b < BranchDeltas.Length; b++)
                BranchDeltas[b] = 0.75f*BranchDeltas[b]
                                  + 0.1f*(float) _random.NextDouble()*(BranchLengthsBest[b] - BranchLengths[b])
                                  + 0.1f*(float) _random.NextDouble()*(BranchLengthGlobalBest[b] - BranchLengths[b]);
        }

        public class ParticleComparer : IComparer<Particle>
        {

            #region IComparer<Particle> Members

            public int Compare(Particle x, Particle y)
            {
                if (x.Likelihood > y.Likelihood)
                    return -1;
                return 1;
            }

            #endregion
        }

    }
}
